from sklearn_crfsuite import CRF

from HMM import read_data


def word2features(sent, i):
    # 提取一个字的特征
    # 因为每个词相邻的词会影响这个词的标记
    # 所以我们使用：
    # 前一个词，当前词，后一个词，
    # 前一个词+当前词， 当前词+后一个词
    # 作为特征
    word = sent[i]
    prev_word = "<s>" if i == 0 else sent[i-1]
    next_word = "</s>" if i == (len(sent)-1) else sent[i+1]

    features = {
        'w': word,
        'w.isnum': word.isdigit(),
        'w-1': prev_word,
        'w-1.isnum': prev_word.isdigit(),
        'w+1': next_word,
        'w+1.isnum': next_word.isdigit(),
        'w-1:w': prev_word+word,
        'w:w+1': word+next_word,
        'bias': 1
    }
    return features


def sent2features(sent):
    # 由序列中每个字的特征组成的列表
    """抽取序列特征"""
    return [word2features(sent, i) for i in range(len(sent))]


class CRFModel(object):
    def __init__(self, algorithm='lbfgs', c1=0.1, c2=0.1, max_iterations=100, all_possible_transitions=False):
        self.model = CRF(algorithm=algorithm,
                         c1=c1,
                         c2=c2,
                         max_iterations=max_iterations,
                         all_possible_transitions=all_possible_transitions)

    def train(self, sentences, tag_lists):
        """训练模型"""
        features = [sent2features(s) for s in sentences]
        self.model.fit(features, tag_lists)

    def test(self, sentences):
        """解码,对给定句子预测其标注"""
        features = [sent2features(s) for s in sentences]
        pred_tag_lists = self.model.predict(features)
        return pred_tag_lists


if __name__ == "__main__":
    train_file = "NER/NER/Chinese/train.txt"
    test_file = "NER/NER/Chinese/validation.txt"
    result_file = "NER/NER/example_data/myCrf.txt"

    train_word_lists, train_tag_lists = read_data(train_file)
    assert len(train_word_lists) == len(train_tag_lists)

    test_word_lists, test_tag_lists = read_data(test_file)
    assert len(test_word_lists) == len(test_tag_lists)

    crf_model = CRFModel()
    crf_model.train(train_word_lists, train_tag_lists)
    results = crf_model.test(test_word_lists)
    assert len(results) == len(test_tag_lists)

    with open(result_file, 'w', encoding='utf-8') as file:
        for word_list, result in zip(test_word_lists, results):
            for i, j in zip(word_list, result):
                file.write(i + " " + j + "\n")
            file.write("\n")
